import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import random
from tqdm import tqdm
from matplotlib.animation import FuncAnimation

class Building:
    """三维建筑结构定义（含随机要素）"""
    def __init__(self, floors=6, size=20):
        self.floors = floors
        self.size = size
        self.stairs = self.generate_random_stairs()
        self.exits = self.generate_random_exits()
        self.obstacles = self.generate_obstacles()
        
    def generate_random_stairs(self):
        stairs = []
        for floor in range(self.floors):
            while True:
                s1 = (random.uniform(2, self.size-2), 
                      random.uniform(2, self.size-2), 
                      float(floor))
                s2 = (random.uniform(2, self.size-2), 
                      random.uniform(2, self.size-2), 
                      float(floor))
                if np.linalg.norm(np.array(s1[:2]) - np.array(s2[:2])) > 5:
                    break
            stairs.extend([s1, s2])
        return stairs
    
    def generate_random_exits(self):
        exits = []
        directions = ['north', 'south', 'east', 'west']
        for _ in range(2):
            dir = random.choice(directions)
            if dir == 'north':
                exits.append((random.uniform(5, self.size-5), self.size-1, 0.0))
            elif dir == 'south':
                exits.append((random.uniform(5, self.size-5), 0.0, 0.0))
            elif dir == 'east':
                exits.append((self.size-1, random.uniform(5, self.size-5), 0.0))
            else:
                exits.append((0.0, random.uniform(5, self.size-5), 0.0))
        return exits
    
    def generate_obstacles(self):
        obstacles = []
        for floor in range(self.floors):
            for _ in range(random.randint(3,5)):
                x = random.uniform(1, self.size-5)
                y = random.uniform(1, self.size-5)
                w = random.uniform(3, 8)
                h = random.uniform(3, 8)
                obstacles.append(('rect', x, y, x+w, y+h, float(floor)))
            for _ in range(random.randint(2,3)):
                cx = random.uniform(5, self.size-5)
                cy = random.uniform(5, self.size-5)
                r = random.uniform(2, 5)
                obstacles.append(('circle', cx, cy, r, float(floor)))
        return obstacles
    
    def is_obstructed(self, pos):
        x, y, z = pos
        for obs in self.obstacles:
            if obs[-1] != z:
                continue
            if obs[0] == 'rect':
                x1, y1, x2, y2 = obs[1:5]
                if x1 <= x <= x2 and y1 <= y <= y2:
                    return True
            elif obs[0] == 'circle':
                cx, cy, r = obs[1:4]
                if (x - cx)**2 + (y - cy)**2 <= r**2:
                    return True
        return False

class Person:
    def __init__(self, id, start_pos):
        self.id = id
        self.pos = np.array(start_pos, dtype=np.float64)
        self.speed = np.clip(np.random.normal(1.0, 0.1), 0.8, 1.2)
        self.total_dist = 0.0
        self.total_time = 0.0
        self.path = []
        self.current_target = None
        
    def update_speed(self):
        a, b, c = 0.71, 0.51, 1.0
        if self.total_time == 0: return
        exponent = (self.total_dist**a)/(self.total_time**b)
        self.speed *= (1 - 0.01*c*np.exp(exponent) + 0.01)
        self.speed = max(self.speed, 0.2)
    
    def find_alternative_path(self, building):
        angle = random.uniform(0, 2*np.pi)
        return self.pos + 2*np.array([np.cos(angle), np.sin(angle), 0])

class Simulation:
    def __init__(self, building, people):
        self.building = building
        self.people = people
        self.time = 0
        
    def generate_path(self, person):
        path = []
        current_pos = person.pos.copy()
        
        while current_pos[2] > 0:
            min_dist = float('inf')
            best_stair = None
            for s in self.building.stairs:
                if s[2] != current_pos[2]:
                    continue
                if self.building.is_obstructed(s):
                    continue
                dist = np.linalg.norm(current_pos[:2] - np.array(s[:2]))
                if dist < min_dist:
                    min_dist = dist
                    best_stair = s
            
            if best_stair is None:
                return None
            path.append(best_stair)
            current_pos = np.array([best_stair[0], best_stair[1], best_stair[2]-1.0])
        
        exit = self.find_safe_exit(current_pos)
        if exit is not None:
            path.append(exit)
        return path
    
    def find_safe_exit(self, pos):
        for exit in self.building.exits:
            if not self.building.is_obstructed(exit):
                return exit
        return None
    
    def move_towards(self, person, target, dt):
        if self.building.is_obstructed(target):
            person.current_target = person.find_alternative_path(self.building)
            return
        
        direction = target - person.pos
        distance = np.linalg.norm(direction)
        step_size = person.speed * dt
        
        if distance <= step_size:
            new_pos = target
        else:
            new_pos = person.pos + (direction/distance)*step_size
        
        if not self.building.is_obstructed(new_pos):
            person.pos = new_pos
            person.total_dist += np.linalg.norm(new_pos - person.pos)
            person.total_time += dt
        else:
            person.current_target = person.find_alternative_path(self.building)
        
        person.update_speed()
    
    def update(self, dt=0.5):
        remaining = []
        for p in self.people:
            if any(np.allclose(p.pos, e, atol=0.5) for e in self.building.exits):
                continue
                
            if not p.path:
                p.path = self.generate_path(p)
                if not p.path:
                    continue
                p.current_target = p.path[0]
            
            if p.current_target is None:
                continue
                
            if np.linalg.norm(p.pos - p.current_target) < 0.5:
                p.path.pop(0)
                if not p.path:
                    p.current_target = None
                    continue
                p.current_target = p.path[0]
            
            self.move_towards(p, p.current_target, dt)
            remaining.append(p)
        
        self.people = remaining
        self.time += dt
        return self.people  # 返回当前人员数据用于可视化更新

def generate_people(building, per_floor=20):
    people = []
    for floor in range(building.floors):
        for _ in range(per_floor):
            while True:
                x = random.uniform(1, building.size-1)
                y = random.uniform(1, building.size-1)
                pos = (x, y, float(floor))
                if not building.is_obstructed(pos):
                    break
            people.append(Person(len(people), pos))
    return people

# 初始化可视化
plt.ion()  # 开启交互模式
fig = plt.figure(figsize=(12, 9))
ax = fig.add_subplot(111, projection='3d')
scatter = None
building = Building(floors=6, size=25)
people = generate_people(building, per_floor=15)
sim = Simulation(building, people)

# 绘制静态元素
def init_visualization():
    ax.clear()
    # 绘制障碍物
    for obs in building.obstacles:
        z = obs[-1]
        if obs[0] == 'rect':
            x1, y1, x2, y2 = obs[1:5]
            xx = [x1, x2, x2, x1, x1]
            yy = [y1, y1, y2, y2, y1]
            ax.plot(xx, yy, [z]*5, c='black', linewidth=1)
        elif obs[0] == 'circle':
            cx, cy, r = obs[1:4]
            theta = np.linspace(0, 2*np.pi, 20)
            xx = cx + r*np.cos(theta)
            yy = cy + r*np.sin(theta)
            ax.plot(xx, yy, [z]*20, c='gray', linewidth=1)
    
    # 绘制楼梯
    stairs = np.array(building.stairs)
    ax.scatter(stairs[:,0], stairs[:,1], stairs[:,2], 
              c='blue', marker='^', s=30, label='Stairs')
    
    # 绘制出口
    exits = np.array(building.exits)
    ax.scatter(exits[:,0], exits[:,1], exits[:,2],
              c='red', marker='X', s=80, label='Exits')
    
    ax.set_xlim(0, building.size)
    ax.set_ylim(0, building.size)
    ax.set_zlim(-0.5, building.floors+0.5)
    ax.set_xlabel('X (m)')
    ax.set_ylabel('Y (m)')
    ax.set_zlabel('Floor')
    plt.title('3D Evacuation Simulation')
    return fig,

# 初始化绘图
fig, = init_visualization()
plt.pause(0.1)

# 动态更新函数
def update_frame(frame):
    global scatter
    people = sim.update(dt=0.5)
    
    if people:
        pos = np.array([p.pos for p in people])
        # 更新散点图数据
        if scatter is None:
            scatter = ax.scatter(pos[:,0], pos[:,1], pos[:,2], 
                                c=pos[:,2], cmap='viridis', s=20)
        else:
            scatter._offsets3d = (pos[:,0], pos[:,1], pos[:,2])
            scatter.set_array(pos[:,2])
    
    ax.set_title(f'Time: {sim.time:.1f}s  Remaining: {len(people)}')
    plt.draw()
    return scatter,

# 创建动画
ani = FuncAnimation(fig, update_frame, frames=200,
                    init_func=init_visualization, blit=False, interval=50)

plt.show(block=True)  # 保持窗口打开